"""
Naive bayesian algorithm.
"""
import numpy as np
from functools import reduce
from Dataset.utils import *

class NB(object):
    def __init__(self, labels=[0,1,2,3]):
        self.labels = labels
        self.vocabulary = None
        self.log_prior = dict([(label, 0) for label in labels])
        self.log_prob = dict([(label, {}) for label in labels])
        
        self.flag = 0
        
    def train(self, train_data):
        """
        Naive bayes training algorithm.
        
        Args:
        - train_data: :numpy.array: (n, 2) the first is the lyric and the second is the label.
        """
        
        # Compute the vocabulary
        print('Now we are compute the vocabulary...')
        corpus, total_corpus = get_corpus(train_data, self.labels)
        self.corpus = corpus
        self.vocabulary, _ = get_vocabulary(total_corpus)
        
        # Compute prior and conditional probabilities
        print('Now we are compute the prior and conditional probabilities...')
        for label in self.labels:
            self.log_prior[label] = np.log(len(corpus[label])) - np.log(len(total_corpus))
            text_c = reduce(lambda x, y: x+y, corpus[label])
            n_c = len(text_c)
            for w in self.vocabulary:
                n_ck = text_c.count(w)
                self.log_prob[label][w] = np.log((n_ck + 1) / (n_c + len(total_corpus)))
        
        # Complete the algorithm
        print('Aha! We finish the training process.')
        self.flag = 1
        
        
    def test(self,test_data):
        """
        Naive bayes test algorithm.
        
        Args:
        - test_data: :numpy.array: (n, 2) the first is the lyric and the second is the label.
        
        Returns:
        - pred_y: :numpy.array: the predict labels.
        - true_y: :numpy.array: the ground truth.
        """
        assert self.flag == 1, 'You have not trained the Naive Bayes algorithm yet.'
        
        true_y = test_data[:,1].astype('int')
        pred_y = np.zeros_like(true_y)
        for i in range(len(test_data)):
            lyric = test_data[i, 0]
            
            pd = np.zeros(len(self.labels))
            # add the prior
            for l in range(len(self.labels)):
                pd[l] += self.log_prior[self.labels[l]]
                
            # add the conditional
            for w in lyric:
                for l in range(len(self.labels)):
                    if w in self.log_prob[self.labels[l]]:
                        pd[l] += self.log_prob[self.labels[l]][w]
            
            # choose the largest one
            pred_y[i] = self.labels[pd.argmax()]
        
        return pred_y, true_y
        
        
        